import time
import psutil
import threading
import os
from functools import wraps

def cpu_monitor(interval=1, filename='cpu_usage.log'):
    # remove existing log file
    if os.path.exists(filename):
        os.remove(filename)

    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # Event to signal the monitoring thread to stop
            stop_event = threading.Event()


            def log_cpu_usage():
                with open(filename, 'a') as f:
                    while not stop_event.is_set():
                        cpu_percent = psutil.cpu_percent(interval=None)
                        memory = psutil.virtual_memory()
                        timestamp = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
                        f.write(f"{timestamp}, {cpu_percent},  {memory.used / (1024 * 1024)}, {memory.total / (1024 * 1024)}\n")
                        f.flush()
                        if stop_event.wait(interval):
                            break

            # Create the monitoring thread
            t = threading.Thread(target=log_cpu_usage)
            t.start()

            # Run the original function
            try:
                t_start = time.time()
                result = func(*args, **kwargs)
                t_end = time.time()
                print(f"[INFO] Running model cost time : {t_end - t_start}s")
            finally:
                # Stop the monitoring thread
                stop_event.set()
                t.join()

        return wrapper
    return decorator